/*
 * MORE
 * Jakubczak, Jennings
 * Copyright (c) 2006 Massachusetts Institute of Technology
 */
#include <click/config.h>
#include <click/confparse.hh>
#include <click/error.hh>
#include <click/glue.hh>
#include <click/straccum.hh>
#include <clicknet/ether.h>
#include <clicknet/llc.h>
#include <clicknet/wifi.h>
#include "more.hh"

CLICK_DECLS

template<typename T>
static inline T max(T a, T b) { return a > b ? a : b; }

/** ****************************************************
 * The MORE element.
 *******************************************************/
int MORE::configure(Vector<String> & conf, ErrorHandler *errh) {
    int ret;

    _data_eth_type = 0xa502;
    _ack_eth_type = 0xa503;
    _log = 0;
    String log_name;

    ret = cp_va_parse(conf, this, errh,
        cpKeywords,
        "ETH", cpEtherAddress, "self MAC address", &_eth,
        "ETHTYPE", cpUnsigned, "Ethernet type", &_data_eth_type,
        "ACKTYPE", cpUnsigned, "Ack Ethernet type", &_ack_eth_type,
        "LOG", cpString, "Log filename", &log_name,
        cpEnd);

    _data_eth_type = htons(_data_eth_type);
    _ack_eth_type = htons(_ack_eth_type);

    if(log_name) {
        _log = fopen(log_name.c_str(), "w");
        if(!_log) {
            return errh->error("Could not open %s", log_name.c_str());
        }
    }

    return ret;
}

#define MULTICAST

void MORE::output_batch(MOREBatch *batch) {
    if(noutputs() >= 3) {
        for(int i = 0; i < batch->D().k(); ++i) {
#ifdef MULTICAST
            const CodedBlock *cb = batch->M()[i];
            // XXX does not really decode for multicast, need to run second run of decoding
            // but this is necessary or otherwise a sink could not forward
#else
            const CodedBlock *cb = batch->D()[i];
#endif
            assert(cb);
            output(2).push(cb->sink()); // this takes time
        }
    }
}

bool MORE::dst_processDATA(MOREFlow &fs, Packet *p, MOREBatch *batch){
    assert(fs.dst >= 0);
    // at destination
#ifdef MULTICAST
    if(!batch->forward(p)) { // this way we can still encode as forwarders
#else
    if(!batch->decode(p)) {
#endif
        return false;
    }
    // NOTE we need to log it here, because we could destroy the batch
    log_data(batch, true, true);

    if(batch->full()) {
        log("Decoded "); if(DEBUG) click_chatter("DECODED!\n");
        output_batch(batch); 

        if(fs.update_dstmap(~(1 << fs.dst))){
            fs.ackflag = PENDING;
            assert(fs.ack_rate > 0);
        }
    }
    return true;
}

// process a MORE encoded packet
/// NOTE: the packet is killed in here
void MORE::processDATA(MOREFlow &fs, Packet *p) {
    bool useful = false;
    const more_pkt *pkt = (const more_pkt *)(p->data());
    const more_data *data = (const more_data *)(pkt+1);

    more_block *mb = (more_block *)(data+1);

    MOREBatch *batch;
    if(fs.src) {
        batch = fs.lookup(pkt->_batch);
    } else {
        if(fs.bcfg.rate < 0) {
            p->kill();
            log("- PRUNED ");
            return;
        }

        // lookup OR create the batch
        fs.bcfg.k = mb->_k;
        fs.bcfg.datasize = mb->_datasize;
        batch = fs.lookup(pkt->_batch, true);
    }
    assert(batch);

    p->pull(sizeof(more_pkt) + sizeof(more_data));
    bool teacher = (data->_metric > batch->c().metric);

    CodedBlock logcb(p, false); // a copy of coeffs just for logs

    if(fs.dst >= 0){
        useful = dst_processDATA(fs, p, batch);
    } else {
        // at forwarder
        if(batch->forward(p)) {
            useful = true;
        }
    }

    if(teacher) {
        batch->_credit += fs.creditf;
    }
    if(!useful || fs.dst < 0) {
        // otherwise we logged in dest_process
        log_data(batch, teacher, useful);
    }
    // log the received packet
    if(DEBUG_CB)
        log(logcb.print());
}

/** **************
 *   push/pull
 *****************/
void MORE::push(int port, Packet *pp) {
    if(_one.stopped) {
        pp->kill();
        return;
    }

    // aah, the heck with it! let's uniquiefy here
    if(pp->shared()) {
        click_chatter("Warning: uniqueify in push() is expensive!\n");
    }
    WritablePacket *p = pp->uniqueify();

    log_start();

    bool tokill = true;
    // process headers
    const click_ether *eh = (const click_ether *) p->data();

    if(!(eh->ether_type == _data_eth_type ||
         eh->ether_type == _ack_eth_type)) {
        click_chatter("Warning: unrecognized eth types received!\n");
        p->kill();
        return;
    }
    bool forMe = (_eth == EtherAddress(eh->ether_dhost));

    switch(port) {
        case 0: log("RX "); break;
        case 1: log("ll "); break;
        case 2: log("sg "); break;
    }
    log_pkt(eh);

    p->pull(sizeof(click_ether)); // maybe we should do that in sub-processing
    const more_pkt *mp = (const more_pkt *)(eh+1);

    // XXX: no flow-state lookup
    MOREFlow &fs = _one;

    // old batches MUST not be serviced!
    if(mp->_batch >= fs.seqno()) {
      switch(port) {
        case 0: { //received
            if(mp->_batch > fs.seqno()) {
                // update yerself
                fs.flush(mp->_batch);
            }
            // it's always good to update dstmap
            // ...especially after a flush
            if(fs.update_dstmap(mp->_dstmap)) {
                fs.source_check();
                // check if eligible to ack
                if(fs.ack_rate > 0
                   && eh->ether_type == _ack_eth_type) {
                    // we should propagate it
                    fs.ackflag = PENDING;
                }
            }

            // any packet of same dstmap from ack_next flushes ackflag
            // NOTE: shortcut support would be good but would require more coordination
            if(fs.ackflag != NONE && fs.dstmap() == mp->_dstmap) {
                if(fs.ack_next == EtherAddress(eh->ether_shost))
                    fs.ackflag = NONE;
            }

            if(eh->ether_type == _ack_eth_type) { // ack
                // NOTE: this is just for debug, ackflag already updated
                if(DEBUG_ACK)
                    click_chatter("received ack for batch %d q %d ackflag %d\n", mp->_batch, fs.queued(), fs.ackflag );
                if(!fs.src && forMe) {
                    // this is just a sanity check
                    if(fs.ack_rate <= 0) {
                        click_chatter("ACK routing misconfigured!\n");
                        break;
                    }
                }
            } else if(eh->ether_type == _data_eth_type) { // data
                tokill = false;
                if(fs.dstmap() > 0) {
                    processDATA(fs, p);
                }
            } else { // unexpected!
                if(DEBUG) {
                    click_chatter("ETH type %x", eh->ether_type);
                    assert(0);
                }
            }
            break;
        }
        case 1: { //txreceipts
            assert(_eth == EtherAddress(eh->ether_shost));

            // NOTE: no point checking rate or ack_rate now
            const click_wifi_extra *ceha = (const click_wifi_extra *) p->all_user_anno();
            assert(ceha->magic == WIFI_EXTRA_MAGIC);
            assert(ceha->flags & WIFI_EXTRA_TX);

            log((int)ceha->retries);

            if(!(WIFI_EXTRA_TX_FAIL & ceha->flags)) {
                log(" s ");
                // The packet was received
                if(eh->ether_type == _ack_eth_type){ // ack
                    if(DEBUG_ACK)
                        click_chatter("ack llacked for seqno %d\n", fs.seqno());
                    if(fs.ackflag == SENT 
                       && mp->_batch == fs.seqno() 
                       && mp->_dstmap == fs.dstmap()) {
                        // NOTE: only SENT is valid here, if PENDING it means it was updated in the meanwhile
                        //ack was received
                        fs.ackflag = NONE;
                    }
                }
            } else {
                log(" f ");
                // the packet was not received
                if(eh->ether_type == _ack_eth_type){
                    //resend ack
                    if(DEBUG_ACK)
                        click_chatter("ack not llacked for seqno %d\n", fs.seqno());
                    if(fs.ackflag == SENT) {
                        // if NONE then we already heard ack_next is up-to-date
                        fs.ackflag = PENDING;
                    }
                }
            }
            break;
        }
        case 2: {
            // signal from ToDevice
            if(PRE_ENCODE) {
                if(eh->ether_type == _data_eth_type){ // data
                    MOREBatch *batch = fs.lookup(mp->_batch);
                    if(batch) {
                        // pass p so that we don't need to allocate again
                        // total #new <= K+1 per batch
                        tokill = false;
                        batch->prepare(p);
                    }
                }
            }
        }
      }
    } else {
      if(DEBUG)
          click_chatter("OLD! batch %d < cumack %d\n", mp->_batch, fs.seqno());
      log("- OLD ");
    }


    if(tokill)
        p->kill();

    log_end();
}

Packet *MORE::pull(int port) {
    WritablePacket *p = 0;
    MOREFlow &fs = _one;

    log_start();

    uint16_t eth_type; // for the packet we create

    //NOTE: port 1 gets higher preference by priosched than port 0.

    if(port == 0) {
        /* data packets */
        if((p = fs.encode())) {
            eth_type = _data_eth_type;
        }
    } else if(port == 1) {
        /* control packets */
        if((p = fs.make_ack())) {
            eth_type = _ack_eth_type;
        } 
    }

    if(p) {
        // complete the ethernet header
        click_ether *eh = (click_ether *)p->data();
        memcpy(eh->ether_shost, _eth.data(), 6);
        eh->ether_type = eth_type;

        log("TX ");
        log_pkt(eh);

        if(eth_type == _data_eth_type) {
            // log the transmitted block
            if(DEBUG_CB) {
                p->pull(sizeof(click_ether) + sizeof(more_pkt) + sizeof(more_data));
                log(CodedBlock(p, false).print());
                p->push(sizeof(click_ether) + sizeof(more_pkt) + sizeof(more_data));
            }
        }

        log_end();
    }
    return p;
}

/** **************
 *   LOGGING
 *****************/
void MORE::log_pkt(const click_ether* eth) {
    if(!_log)
        return;

    // NOTE: for the case of a and d, we really want to log dest
    // from ether address
    bool self = (EtherAddress(eth->ether_shost) == _eth);
    if(self)
        _log_sa << EtherAddress(eth->ether_dhost) << " ";
    else
        _log_sa << EtherAddress(eth->ether_shost) << " ";

    const more_pkt *mp = (const more_pkt *)(eth+1);

    // batch no
    _log_sa << mp->_batch << " " << mp->_dstmap << " ";

    uint etype = eth->ether_type;
    if(etype == _data_eth_type) {
        _log_sa << (self ? "d " : "D ");
        _log_sa << ((const more_data *)(mp+1))->_tx << " ";
    } else if(etype == _ack_eth_type) {
        _log_sa << (self ? "a " : "A ");
    } else {
        _log_sa << (self ? "! " : "? ");
    }
}
void MORE::log_data(const MOREBatch *batch, bool teacher, bool useful) {
    if(!_log)
        return;

    _log_sa << (teacher ? "T " : "S ") << (int)useful << " ";

    if(batch)
        _log_sa << batch->rank() << " " << 0 << " " << batch->M().rank() << " " << batch->_credit;
}
void MORE::log_end() {
    if(!_log)
        return;

    timeval finish;
    click_gettimeofday(&finish);
    timersub(&finish, &_log_start, &finish);

    StringAccum sa;
    sa << _log_start << " " << finish << " ";

    // NOTE: actual output is here!
    fprintf(_log, "%s %s\n", sa.take_string().c_str(), _log_sa.take_string().c_str());
}


/** *************************
 * ControlSocket handlers
 **************************/
enum {H_START, H_STOP, H_SRC, H_DST, H_SET, H_INFLOW, H_OUTFLOW, H_FLUSH, H_ACK};

int MORE::write_param(const String &in_s, Element *e, void *vparam,
                ErrorHandler *errh)
{
    MORE *f = (MORE *)e;
    String s = cp_uncomment(in_s);
    switch((int)vparam) {
        case H_START: {
            MOREFlow &fs = f->_one;
            if(! fs.src){
                return errh->error("Not configured to be src, start forbidden!");
            }
            fs.source_check();
        } break;
        case H_STOP: {
            f->_one.stop();
        } break;
        case H_SET: {
            // DATA forwarding
            // <metric> <rate> <next.mac>
            Vector<String> args;
            cp_spacevec(in_s, args);

            MOREFlow &fs = f->_one;
            assert(fs.seqno() == 0);
            fs.stopped = false;

            if (args.size() != 2) {
                return errh->error("Must have (metric, rate) currently has %d", args.size());
            }
            if (!cp_unsigned(args[0], &fs.bcfg.metric)) {
                return errh->error("Couldn't read metric");
            }
            if (!cp_integer(args[1], &fs.bcfg.rate)) {
                return errh->error("Couldn't read rate");
            }
        } break;
        case H_SRC: {
            // <datasize> <k> <batch_window>
            Vector<String> args;
            cp_spacevec(in_s, args);

            MOREFlow &fs = f->_one;
            fs.src = true;

            if (args.size() != 3) {
                return errh->error("Must have (datasize, k, dstmap) currently has %d", args.size());
            }

            if (!cp_integer(args[0], &fs.bcfg.datasize)) {
                return errh->error("Couldn't read datasize");
            }
            if (!cp_integer(args[1], &fs.bcfg.k)) {
                return errh->error("Couldn't read k");
            }
            if (!cp_unsigned(args[2], &fs.src_dstmap)) {
                return errh->error("Couldn't read dstmap");
            }
        } break;
        case H_DST: {
            MOREFlow &fs = f->_one;
            if(!cp_integer(in_s, &fs.dst)) {
               return errh->error("integer argument required");
            }
        } break;
        case H_OUTFLOW: case H_INFLOW: {
            // add credit info
            Vector<String> args;
            cp_spacevec(in_s, args);

            Vector<float> fargs;

            for(Vector<String>::iterator it = args.begin(); it != args.end(); ++it) {
               double d;
               if(!cp_double(*it, &d)) return errh->error("a list of doubles required");
               fargs.push_back(d);
            }
            if((int)vparam == H_OUTFLOW)
                f->_one.set_outflow(fargs); 
            if((int)vparam == H_INFLOW)
                f->_one.add_inflow(fargs); // at source inflow = sum tx rates of all neighbors
        } break;
        case H_ACK: {
            // ACK forwarding (dst -> src)
            // <ack_rate> <ack_next.mac> ...
            Vector<String> args;
            cp_spacevec(in_s, args);

            MOREFlow &fs = f->_one;
            fs.ackflag = NONE;
            fs.stopped = false;

            if (args.size() != 2) {
                return errh->error("Must have (rate, next)");
            }
            if (!cp_integer(args[0], &fs.ack_rate)) {
                return errh->error("Couldn't read rate");
            }
            if (!cp_ethernet_address(args[1], &fs.ack_next)) {
                return errh->error("Couldn't read next hop");
            }
        } break;
        case H_FLUSH:
            f->flush_log();
            break;

        default: break;
    }
    return 0;
}


void MORE::add_handlers() {
    add_write_handler("start", write_param, (void *)H_START);
    add_write_handler("stop", write_param, (void *)H_STOP);
    add_write_handler("set", write_param, (void *)H_SET);
    add_write_handler("src", write_param, (void *)H_SRC);
    add_write_handler("dst", write_param, (void *)H_DST);
    add_write_handler("outflow", write_param, (void *)H_OUTFLOW);
    add_write_handler("inflow", write_param, (void *)H_INFLOW);
    add_write_handler("ack", write_param, (void *)H_ACK);
    add_write_handler("flush", write_param, (void *)H_FLUSH);
}


#include <click/bighashmap.cc>
#include <click/hashmap.cc>
#include <click/vector.cc>
#include <click/dequeue.cc>
// template class declarations here

EXPORT_ELEMENT(MORE)
CLICK_ENDDECLS
